import zhipuai
import time

zhipuai.api_key = "e07484c365a3d8b468a4d9e06b36317a.momrwta4EftapAwx"


def invoke_example():
    response = zhipuai.model_api.invoke(
        model="chatglm_turbo",
        prompt=[{"role": "user", "content": "人工智能"}],
        top_p=0.7,
        temperature=0.9,
    )
    print(response['data']['choices'][0]['content'])

def async_invoke_example():
    response = zhipuai.model_api.async_invoke(
        model="chatglm_turbo",
        prompt=[{"role": "user", "content": "人工智能"}],
        top_p=0.7,
        temperature=0.9,
    )
    print(response)
    if response['code'] == 200:
        return response['data']['task_id']
    if response['code'] != 1302:
        return None  

def sse_invoke_example():
    response = zhipuai.model_api.sse_invoke(
        model="chatglm_turbo",
        prompt=[{"role": "user", "content": "人工智能"}],
        top_p=0.7,
        temperature=0.9,
    )

    for event in response.events():
        if event.event == "add":
            print(event.data)
        elif event.event == "error" or event.event == "interrupted":
            print(event.data)
        elif event.event == "finish":
            print(event.data)
            print(event.meta)
        else:
            print(event.data)

def query_async_invoke_result_example(task_id):
    response = zhipuai.model_api.query_async_invoke_result(task_id)
    if response['success']:
        if response['data']['task_status'] == 'SUCCESS':
            return {'status': 1, 'content': response['data']['choices'][0]['content']}
        else:
            return {'status':0}
    else:
        return {'status': -1}
    
    
#invoke_example()

'''
task_id = async_invoke_example()
while True:
    resp = query_async_invoke_result_example(task_id)
    print(resp)
    if resp['status'] != 0:
        break
    time.sleep(5)
'''
sse_invoke_example()
